from random import random
from elementtree.ElementTree import parse
from blitz import *
from math import ceil

try:
   import psyco
   psyco.full()
except:
   pass

class Tabletop:
   def __init__(self):
      self.units = []
      return
   
   def addUnit(self, unit, pilot, name=None):
      """Adds a unit to the table"""
      self.units.append(Unit(unit, pilot))
      if name:
         self.units[-1].id = name
      else:
         self.units[-1].id = str(random())
      return self.units[-1]

   def getDistance(self, posA, posB):
      """Gets the distance between two points"""
      x = posA[0] - posB[0]
      y = posA[1] - posB[1]
      if x == 0:
         distance = y
      elif y == 0:
         distance = x
      else:
         distance = (y**2 + x**2)**.5
      if distance < 0:
         distance *= -1
      return ceil(distance)

   def getTargets(self, radius, pos):
      """Returns all targets in radius distance from pos"""
      targets = []
      for unit in self.units:
         if self.getDistance(pos, unit.pos) <= radius:
            targets.append(unit)
      return targets
      
   def getScatter(self, distance, pos):
      """Gives a point in a random direction distance from pos"""
      xdir = ceil(random() * 3) - 2
      ydir = ceil(random() * 3) - 2
      if xdir != 0 and ydir != 0:
         distance *= 0.5
      x = pos[0] + (distance * xdir)
      y = pos[1] + (distance * ydir)
      return x, y

   def endTurn(self):
      """Starts the next turn"""
      for unit in self.units:
         unit.untap()

   def getUnit(self, name):
      """Selects a unit by its name"""
      for unit in self.units:
         if unit.name == name:
            return unit

def testWeapons(table, attacker, defender, distance, weapons, tests=1000, stopAtPenelty=-10):
   """Tests the effectiveness of a weapon or weapons against a single target at a given range"""
   attacker.pos = [0, 0]
   defender.pos = [0, distance]
   
   # This figures out how best to apply IF, RoF and AE.
   firingMode = {}
   for weapon in weapons:
      weapon = attacker.weapons[weapon]
      indirect = 0
      ae = 0
      rof = weapon.rof
      spray = 0
      burst = 0
      
      if not weapon.melee:
         atkMod = attacker.getRangedAttackMod(table, defender.pos, weapon)
         if weapon.indirect:
            if atkMod < 0:  # Sets the weapon to use indirect if there's a negative attack modifier. Is that best?
               indirect = 1
               # Checks to use AE. It would be better if this were done every shot but eh.
               if weapon.ae:
                  if weapon.ai:
                     ae = 1
                  elif distance > weapon.ae:
                     if defender.defSpeedMod[defender.movementMode][defender.speedband] > 0:
                        if (weapon.dm - weapon.ae) * 2 > defender.armor:
                           ae = 1                    
                        elif (weapon.dm - weapon.ae + rof) * 2 > defender.armor:
                           spray = weapon.dm - weapon.ae + weapon.rof - defender.armor
                           if weapon.ae + spray > distance:
                              spray = distance - weapon.ae
                           burst = weapon.rof - spray
                           rof = 0
                     else:
                        if weapon.dm - weapon.ae > defender.armor:
                           ae = 1
                        elif weapon.dm + rof - weapon.ae > defender.armor:
                           ae = 1
                           spray = weapon.dm - weapon.ae + rof - defender.armor
                           if weapon.ae + spray > distance:
                              spray = distance - weapon.ae
                           burst = weapon.rof - spray
                           rof = 0
               if rof:
                  if weapon.dm >= defender.armor:
                     spray = rof
                     if spray > distance:
                        spray -= distance
                        burst = weapon.rof - spray
                     rof = 0
                  elif weapon.dm + rof >= defender.armor:
                     spray = weapon.dm + rof - defender.armor
                     if spray > distance:
                        spray -= distance
                     burst = weapon.rof - spray
                     rof = 0
                  else:
                     spray = rof
                     rof = 0
            else:
               if weapon.ae:
                  if weapon.ai:
                     ae = 1
               elif distance > weapon.ae:
                  if defender.defSpeedMod[defender.movementMode][defender.speedband] > 0:
                     if (weapon.dm - weapon.ae + rof) * 2 > defender.armor:
                        ae = 1
               burst = rof
               rof = 0
         else:
            if weapon.ae:
               if weapon.ai:
                  ae = 1
            elif distance > weapon.ae:
               if defender.defSpeedMod[defender.movementMode][defender.speedband] > 0:
                  if (weapon.dm - weapon.ae + rof) * 2 > defender.armor:
                     ae = 1
            burst = rof
            rof = 0
      firingMode[weapon.code] = spray, burst, ae, indirect
   
   total = 0
   for i in xrange(0, tests):
      turns = 0
      defender.boxes = defender.defaultBoxes
      defender.stunned = 0
      defender.penelty = 0
      while not defender.isDead() and defender.penelty > stopAtPenelty and turns < 20:
         turns += 1
         total += 1
         for weapon in weapons:
            if attacker.weapons[weapon].melee and distance == 1:
               attacker.meleeAttack(table, defender.pos, weapon)
            elif distance <= attacker.weapons[weapon].rng[-1]:
               damage = attacker.rangedAttack(table, defender.pos, weapon, spray=firingMode[weapon][0], burst=firingMode[weapon][1], ae=firingMode[weapon][2], indirect=firingMode[weapon][3])
   return total / float(tests)


def testUnit(attacker, defender, atkPilot=pilot, defPilot=pilot, options=0, swaps=0, plot=0):
   table = Tabletop()
   attacker = table.addUnit(attacker, atkPilot)
   attacker.speedband = 1
   if options:
      for option in options:
         attacker.addWeapon(option)
   if swaps:
      for swap in swaps:
         attacker.removeWeapon(swap)
         
   defender = table.addUnit(defender, defPilot)
   defender.speedband = 1
   defender.movementMode = 'Ground'

   # Tests each weapon the attacker has at all of its rangebands.
   results = {}
   for weapon in attacker.weapons:
      results[weapon] = []
      for distance in attacker.weapons[weapon].rng:
         score = testWeapons(table, attacker, defender, distance, [weapon])
         results[weapon].append(score)

   if attacker.actions > 1:
      # Multiaction units are still done bluntly but it's less important as they kill things faster.
      # This sorts the unit's weapons and test the top (number of action) ones at all ranges 1-24.
      lowestResults = []
      for distance in xrange(1, 25):
         bestWeapons = {}
         for weapon in results:
            for rangeband in attacker.weapons[weapon].rng:
               if distance <= rangeband:
                  bestWeapons[weapon] = results[weapon][attacker.weapons[weapon].getRangeband(rangeband)-1]
         items = bestWeapons.items()
         backitems = [[v[1], v[0]] for v in items]
         backitems.sort()
         sortedWeapons = [backitems[i][1] for i in range(0, len(backitems))]
         lowestResults.append(testWeapons(table, attacker, defender, distance, sortedWeapons[:attacker.actions]))
   else:
      # This converts the breakpoint scores of results into a full list of the best scores from ranges 1-24
      lowestResults = range(20, 44)
      for distance in xrange(1, 25):
         for weapon in results:
            for rangeband in attacker.weapons[weapon].rng:
               if distance <= rangeband:
                  if results[weapon][attacker.weapons[weapon].getRangeband(rangeband)-1] < lowestResults[distance-1]:
                     lowestResults[distance-1] = results[weapon][attacker.weapons[weapon].getRangeband(rangeband)-1]
   if plot:
      try:
         from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
         from matplotlib.figure import Figure
         option = ''
         if options:
            for i in options:
               option += i + '/'
         option = option[:-1]
         fig = Figure()
         canvas = FigureCanvas(fig)
         ax = fig.add_subplot(111)
         ax.set_title('%s %s(%s/%s) - %s(%s/%s)' % (option, attacker.name, atkPilot['atk'], atkPilot['def'], defender.name, defPilot['atk'], defPilot['def']))
         ax.set_ylabel('Shots to kill')
         ax.set_xlabel('Distance')
         ax.set_yticks(range(0, 21, 2))
         ax.set_xticks(range(0, 25, 2))
         ax.grid (True)   
         ax.plot(xrange(1, 25), lowestResults)
         ax.axis([1, 24, 0, 20])
         canvas.print_figure('%s %s(%s%s) - %s(%s%s).png' % (option, attacker.name, atkPilot['atk'], atkPilot['def'], defender.name, defPilot['atk'], defPilot['def']))
      except:
         print 'Install matplotlib'
         
   return results, lowestResults

def compareUnits(unitA, unitB, defender, pilotA=pilot, pilotB=pilot, defPilot=pilot, optionAs=0, optionBs=0, aSwaps=0, bSwaps=0, plot=0):
   """Compares the difference between two units against a specific target. If plot, it will make three graphs. One
      For unitA vs defender, one for unitB vs defender, and one of unitA and unitB vs defender"""
   unitAResults = testUnit(unitA, defender, atkPilot=pilotA, defPilot=defPilot, plot=plot, options=optionAs, swaps=aSwaps)
   unitBResults = testUnit(unitB, defender, atkPilot=pilotB, defPilot=defPilot, plot=plot, options=optionBs, swaps=bSwaps)

   if plot:
      try:
         from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
         from matplotlib.figure import Figure
         optionA = ''
         if optionAs:
            for i in optionAs:
               optionA += i + '/'
            optionA = optionA[:-1] + ' '
         optionB = ''
         if optionBs:
            for i in optionBs:
               optionB += i + '/'
            optionB = optionB[:-1] + ' '
         fig = Figure()
         canvas = FigureCanvas(fig)
         ax = fig.add_subplot(111)
         values = optionA, unitA, pilotA['atk'], optionB, unitB, pilotB['atk'], defender, defPilot['def']
         ax.set_title('%s%s(%s) + %s%s(%s) - %s(%s)' % values)
         ax.set_ylabel('Shots to kill')
         ax.set_xlabel('Distance')
         ax.set_yticks(range(0, 21, 2))
         ax.set_xticks(range(0, 25, 2))
         ax.grid (True)   
         ax.plot(xrange(1, 25), unitAResults[1], label=unitA)
         ax.plot(xrange(1, 25), unitBResults[1], label=unitB)
         ax.legend(loc=2)
         ax.axis([1, 24, 0, 20])
         canvas.print_figure('%s %s(%s) + %s %s(%s) - %s(%s).png' % values)
      except:
         print 'Install matplotlib'

   return compareResults(unitAResults[1], unitBResults[1])
   

def compareResults(resultsA, resultsB):
   """ Compares the lowestResults of two units and gives you the percentage difference of A from B."""
   results = []
   for distance in xrange(0, 23):
      results.append(100.0/resultsA[distance]*resultsB[distance])
   return results